import json
import os
import random
import os,shutil


root_path = r'D:\工作相关内容\公司项目\禁限运危化品货源\train_data'
num_map_dict = {
    "normal_cargo":4000,
    "bad_cargo":5000,
    "deal_cargo":2000,
}


def getfiles(fold_name,return_fold=False):
    file_lists = []
    for file in os.listdir(fold_name):
        real_path = os.path.join(fold_name, file)
        if not return_fold  and os.path.isfile(real_path):
            file_lists.append(real_path)
        elif return_fold and os.path.isdir(real_path):
            file_lists.append(real_path)
    return file_lists

def get_samplesfiles_from_fold(fold_name,sample_nums=-1):
    file_lists = getfiles(fold_name)
    if len(file_lists) > sample_nums:
        sample_list = random.choices(file_lists,k=sample_nums)
    else:
        sample_list = file_lists

    if sample_nums == -1:
        sample_list = file_lists
    left_list = set(file_lists) - set(sample_list)
    return sample_list,list(left_list)


def sample_fold_files(fold_name,sample_num,sample_all=False):
    result = []
    folds_list = getfiles(fold_name,return_fold=True)
    if not sample_all:
        if len(folds_list) >= sample_num:
            sample_folds = random.choices(folds_list,k=sample_num)
            for fold in sample_folds:
                result.extend(get_samplesfiles_from_fold(fold,1)[0])
        else:
            left = []
            sample_each_fold = int(sample_num/len(folds_list))
            for fold in folds_list:
                sample_files,left_files = get_samplesfiles_from_fold(fold,sample_each_fold)
                result.extend(sample_files)
                left.extend(left_files)

            left_need_sample_num = sample_num - len(result)
            result.extend(random.sample(left,k=left_need_sample_num))
    else:
        for fold in folds_list:
            sample_files, left_files = get_samplesfiles_from_fold(fold, -1)
            result.extend(sample_files)


    return result

def generator_docanno_file():
    index = 0
    save_root = r'D:\工作相关内容\公司项目\禁限运危化品货源\train_data\doccano\train1.json'
    output = []
    with open(save_root,'w',encoding='utf-8') as fwrite:
        for key,value in num_map_dict.items():

            check_path = os.path.join(root_path,key)
            generator_files = sample_fold_files(check_path,value)

            print(len(generator_files))
            for file in generator_files:
                with open(file,encoding='utf-8') as fread:
                    try:
                        text = fread.read()
                        item = {}
                        item["text"] = text
                        output.append(item)
                    except Exception as e:
                        print("error:"+file)
        fwrite.write(json.dumps(output,indent=4,ensure_ascii=False))



def generator_badnamedocanno_file():
    index = 0
    save_root = r'D:\工作相关内容\公司项目\禁限运危化品货源\train_data\doccano\badname.json'
    output = []
    with open(save_root,'w',encoding='utf-8') as fwrite:
        check_path = os.path.join(root_path,"bad_cargo")
        generator_files = sample_fold_files(check_path,5000)

        print(len(generator_files))
        for file in generator_files:
            with open(file,encoding='utf-8') as fread:
                try:
                    text = fread.read()
                    item = {}
                    item["text"] = text
                    output.append(item)
                except Exception as e:
                    print("error:"+file)
        random.shuffle(output)
        fwrite.write(json.dumps(output))
















def count_files(fold_path):
    file_count = 0
    for foldname in os.listdir(fold_path):
        real_path = os.path.join(fold_path,foldname)
        if not os.path.isdir(real_path):
            file_count += 1
        else:
            file_count += count_files(real_path)
    return  file_count

if __name__ == '__main__':
    from dangerour_word.utils_dangerous_words import generator_docanno_file
    data_fold = r'D:\工作相关内容\公司项目\禁限运危化品货源\20220106\train'
    save_path = r'D:\工作相关内容\公司项目\禁限运危化品货源\20220106\doccano.json'
    generator_docanno_file(data_fold,save_path)


